package main

import (
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net"
	"net/http"
	"time"
)

//var upgrader = websocket.Upgrader{}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有来源
	},
}

func main() {
	http.HandleFunc("/", handleWebSocket)
	err := http.ListenAndServe(":18182", nil)
	if err != nil {
		log.Fatal("Error starting the server: ", err)
	}
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Error upgrading to WebSocket:", err)
		return
	}
	defer conn.Close()

	fmt.Println("New WebSocket connection established from:", conn.RemoteAddr().String())

	// 设置读取消息的超时时间为5秒
	conn.SetReadDeadline(time.Now().Add(5 * time.Second))

	// 循环读取客户端发送的消息
	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			// 处理错误
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Println("WebSocket connection closed:", err)
			} else if err, ok := err.(net.Error); ok && err.Timeout() {
				log.Println("Read timeout occurred.")
			} else {
				log.Println("Error reading message:", err)
			}
			// 继续等待下一次超时
			conn.SetReadDeadline(time.Now().Add(5 * time.Second))
			continue
		}

		handleMessage(conn, message)
		// 如果收到特定的结束标识，可以使用 break 退出循环，关闭连接
		if string(message) == "quit" {
			break
		}

		// 重置超时时间为5秒
		conn.SetReadDeadline(time.Now().Add(5 * time.Second))
	}

	fmt.Println("WebSocket connection closed for:", conn.RemoteAddr().String())
}

func handleMessage(conn *websocket.Conn, message []byte) {
	currentTime := time.Now()
	timeString := currentTime.Format("2006-01-02 15:04:05")

	fmt.Printf("%v: Received message: %x\n", timeString, message)

	// 发送消息回客户端
	if string(message) == "hello" {
		response := []byte("Hello from server!")
		err := conn.WriteMessage(websocket.TextMessage, response)
		if err != nil {
			log.Println("Error sending message:", err)
		}
	}
}
